package com.hgw.AiToImage;

import cloud.liblibai.client.LibLib;
import cloud.liblibai.openapi.client.ApiClient;
import cloud.liblibai.openapi.client.JSON;
import cloud.liblibai.openapi.client.model.*;
import com.google.gson.*;
import com.hgw.AiToImage.config.ConfigLoader;
import spark.Request;
import spark.Response;
import spark.Route;
import spark.Spark;

import java.lang.reflect.Field;
import java.lang.reflect.Type;

import static spark.Spark.get;

public class liblibAIToImage {
    // 用于接收前端请求的DTO
    static class PromptRequest {
        private String prompt;

        public String getPrompt() {
            return prompt;
        }
    }

    public static void main(String[] args) {
        // 设置服务器端口
        Spark.port(8080);

        // 允许跨域请求
        Spark.before((request, response) -> {
            response.header("Access-Control-Allow-Origin", "*");
            response.header("Access-Control-Allow-Methods", "POST");
            response.header("Access-Control-Allow-Headers", "Content-Type");
        });

        // 添加根路径路由 - 返回简单的欢迎信息
        get("/generate", (request, response) -> {
            return "Welcome to LiblibAI Service! Use POST /generate to create images.";
        });

        // 添加 favicon 路由 - 避免 404 错误
        get("/favicon.ico", (request, response) -> {
            response.status(204); // No Content
            return "";
        });

        // 创建图像生成API端点
        Spark.post("/generate", new Route() {
            @Override
            public Object handle(Request request, Response response) {
                try {
                    // 解析请求体
                    Gson gson = new Gson();
                    PromptRequest promptRequest = gson.fromJson(request.body(), PromptRequest.class);
                    // 验证提示词
                    if (promptRequest == null || promptRequest.getPrompt() == null || promptRequest.getPrompt().trim().isEmpty()) {
                        response.status(400);
                        // 返回JSON格式错误
                        response.type("application/json");
                        return "{\"error\": \"Prompt is required\"}";
                    }
                    // 调用图像生成方法
                    String imageUrl = generateImage(promptRequest.getPrompt());
                    // 设置响应类型为JSON
                    response.type("application/json");
                    // 返回JSON格式响应
                    return "{\"imageUrl\": \"" + imageUrl + "\"}";
                } catch (Exception e) {
                    e.printStackTrace();
                    response.status(500);
                    response.type("application/json");
                    return "{\"error\": \"" + e.getMessage().replace("\"", "\\\"") + "\"}";
                }
            }
        });

        System.out.println("LiblibAI 服务已启动，访问地址: http://192.168.60.42:8080/generate");
    }

    private static String generateImage(String prompt) throws Exception {
        String accessKey = ConfigLoader.getAccessKey();
        String secretKey = ConfigLoader.getSecretKey();

        // 创建API客户端
        LibLib api = new LibLib(accessKey, secretKey);
        ApiClient apiClient = api.getApiClient();

        // 创建自定义Gson配置，解决seed字段超出int范围的问题
        Gson customGson = new GsonBuilder()
                .setLenient()
                .serializeNulls()
                .disableHtmlEscaping()
                // 添加自定义类型适配器，处理大数值的seed字段
                .registerTypeAdapter(Integer.class, new JsonDeserializer<Integer>() {
                    @Override
                    public Integer deserialize(JsonElement json, Type typeOfT, JsonDeserializationContext context)
                            throws JsonParseException {
                        try {
                            // 尝试解析为int
                            return json.getAsInt();
                        } catch (Exception e) {
                            // 如果解析失败，尝试解析为long并转换为int（可能丢失精度）
                            try {
                                long value = json.getAsLong();
                                return (int) value;
                            } catch (Exception ex) {
                                // 如果还是失败，返回默认值
                                return 0;
                            }
                        }
                    }
                })
                .create();

        // 使用反射设置自定义Gson
        JSON json = apiClient.getJSON();
        Field gsonField = JSON.class.getDeclaredField("gson");
        gsonField.setAccessible(true);
        gsonField.set(json, customGson);

        // 构建文生图请求（使用传入的提示词）
        TextToImageRequest request = getTextToImageRequest(prompt);

        // 提交异步请求
        System.out.println("正在提交文生图请求...");
        SubmitResponse submitResponse = api.submitTextToImage(request);
        String generateUuid = submitResponse.getData().getGenerateUuid();
        System.out.println("任务ID: " + generateUuid);

        // 轮询查询状态
        System.out.println("正在查询生成状态...");
        StatusResponse status;
        int attempts = 0;
        final int maxAttempts = 20;
        do {
            status = api.getStatus(new StatusRequest().generateUuid(generateUuid));
            System.out.println("当前状态: " + status.getData().getGenerateStatus());
            if (status.getData().getGenerateStatus() != GenerateStatus.SUCCEED) {
                Thread.sleep(5000);
                attempts++;
            }
        } while (status.getData().getGenerateStatus() != GenerateStatus.SUCCEED && attempts < maxAttempts);

        // 获取图片URL
        if (status.getData().getGenerateStatus() == GenerateStatus.SUCCEED) {
            Gson gson = new Gson();
            String jsonData = gson.toJson(status.getData());
            JsonObject dataObj = gson.fromJson(jsonData, JsonObject.class);
            JsonArray imagesArray = dataObj.getAsJsonArray("images");
            // 返回第一个图片URL
            if (imagesArray != null && imagesArray.size() > 0) {
                JsonElement element = imagesArray.get(0);
                if (element.isJsonObject()) {
                    JsonObject imageObj = element.getAsJsonObject();
                    if (imageObj.has("url")) {
                        String url = imageObj.get("url").getAsString();
                        System.out.println("获取到图片URL: " + url);
                        return url;
                    } else if (imageObj.has("imageUrl")) {
                        String imageUrl = imageObj.get("imageUrl").getAsString();
                        System.out.println("获取到图片URL: " + imageUrl);
                        return imageUrl;
                    }
                } else {
                    String url = element.getAsString();
                    System.out.println("获取到图片URL: " + url);
                    return url;
                }
            }
        }

        throw new Exception("Image generation failed or timed out");
    }

    private static TextToImageRequest getTextToImageRequest(String prompt) {
        TextToImageRequest request = new TextToImageRequest();
        request.setTemplateUuid("5d7e67009b344550bc1aa6ccbfa1d7f4");
        TextToImageRequestGenerateParams params = new TextToImageRequestGenerateParams();
        params.setPrompt(prompt); // 使用传入的prompt
        params.setWidth(1024);
        params.setHeight(1024);
        params.setImgCount(1);
        params.setSteps(30);
        request.setGenerateParams(params);
        return request;
    }
}